from rest_framework.throttling import SimpleRateThrottle
from django.core.cache import cache as default_cache


class IPThrottle(SimpleRateThrottle):
    scope = "ip"
    # THROTTLE_RATES = {"XXX": "5/m"}  # 限制访问频率 一分钟访问5次
    cache = default_cache  # redis配置

    # 获取唯一标识，后续可以连接redis做操作
    def get_cache_key(self, request, view):
        ident = self.get_ident(request)  # 获取请求用户IP（去request中找请求头）
        return self.cache_format % {'scope': self.scope, 'ident': ident}


class UserThrottle(SimpleRateThrottle):
    scope = "user"
    # THROTTLE_RATES = {"XXX": "5/m"}  # 限制访问频率 一分钟访问5次
    cache = default_cache  # redis配置

    # 获取唯一标识，后续可以连接redis做操作
    def get_cache_key(self, request, view):
        ident = request.user.pk  # 用户ID
        return self.cache_format % {'scope': self.scope, 'ident': ident}